/*Includes*/
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <errno.h>
#include <assert.h>
#include <sys/types.h>
#include <sys/shm.h>
#include <sys/ipc.h>
#include <sys/mman.h>
#include <fcntl.h>
#include <sys/wait.h>

#include "hw1.h"
/**********/

/*Definitions*/
#define SLEEP_SEC 1
#define SHARED_MEM_FILE "/tmp"
#define SHARED_MEM_ID 'R'
#define SHARED_MEM_FILE_1 "/dev/zero"
#define SHARED_MEM_ID_1 'L'
/*************/

/*MACROS*/
#define INDEX_FATHER_TO_CHILD(index, numOfChildren, childIndex) (index * numOfChildren + childIndex + 1)
#define INDEX_CHILD_TO_FATHER(index, numOfChildren) ((index - 1) / numOfChildren)
#define MAX(num1, num2) (num1 >= num2 ? num1 : num2)
/********/

/*Globals*/
static pid_t* processArr;
static int* isDone;
/*********/

/************************************************************************/
/*							Declarations                                */
/************************************************************************/

/**
 * builds a full process tree.
 * size of new tree will be numOfChildren^(depth+1) - 1
 * @param depth: new tree depth
 * @param numOfChildren: number of children per node
 * @param processArr: process array to fill with pids
 */
void build_process_tree(int depth, int numOfChildren);

/**
 * internal implementation for recursive build of process tree
 * @param depth: new tree depth
 * @param numOfChildren: number of children per node
 * @param processArr: process array to fill with pids
 * @param index: index of the process in the tree which is the root
 * of the sub-tree that will be built underneath it
 */
void do_build_process_tree(int depth, int numOfChildren, int index);

/**
 * calculates the size of a full tree with the supplied depth 
 * and number of children per node
 * @param depth: new tree depth
 * @param numOfChildren: number of children per node
 */
int calc_tree_size(int depth, int numOfChildren);

/**
 * calculates the least common ancestor between index1 and index2
 * in a complete tree
 * @param index1: index of first node in an in-order indexed complete tree
 * @param index2: index of second node in an in-order indexed complete tree
 * @param depth: depth of the tree
 * @param numOfChildren: number of children per node
 */
int calc_lca(int index1, int index2, int depth, int numOfChildren);

/**
 * calculates the path from the root of an in-order indexed complete tree
 * to the specified index
 * @param index1: node index
 * @param numOfChildren: number of children per node
 * @param numOfChildren: output, path array to fill
 */
void path_from_root(int index, int *path, int numOfChildren, int* pathLength);

/**
 * calculates the shortest path from index1 to index2
 * @param index1: index of first node in an in-order indexed complete tree
 * @param index2: index of second node in an in-order indexed complete tree
 * @param depth: depth of the tree
 * @param numOfChildren: number of children per node
 * @param numOfChildren: output, path array to fill
 */
int calc_path(int indexSrc, int indexDst, int depth, int numOfChildren, int* path, int pathLength);

/**
 * calculates x^y
 * @param x: base unsigned integer
 * @param y: exponent unsigned integer
 */
unsigned int pow (unsigned int x,unsigned int y);

/************************************************************************/

/************************************************************************/
/*							Implementation                              */
/************************************************************************/


void build_process_tree(int depth, int numOfChildren) {
	int retVal;
	retVal = fork();
	if (retVal < 0) {
		perror("Failed to create child process\n");
		return;
	}
	if (retVal == 0) {
		processArr[0] = getpid();
		do_build_process_tree(depth, numOfChildren, 0);	
	}
}

void do_build_process_tree(int depth, int numOfChildren, int index) {
	assert(depth >=0 && numOfChildren > 0);
	int retVal;
	if (depth != 0)
	{
		for (int i = 0; i < numOfChildren; i++) {
			retVal = fork();
			if (retVal < 0) /*fail*/ {
				perror("Failed to create child process\n");
				return;
			}
			if (retVal == 0) /*child process*/ {
				index = INDEX_FATHER_TO_CHILD(index, numOfChildren, i);
				processArr[index] = getpid();
				do_build_process_tree(
					depth - 1, 
					numOfChildren, 
					index); //recursive call
				break;
			}
		}
	}
	//printf ("process %d sleeping\n", index);
	fflush(stdout);
	while (*isDone == 0) {
		sleep(SLEEP_SEC);
	}
	//printf ("process %d closing\n", index);
	fflush(stdout);
	int status;
	while(wait(&status) != -1) {}
	exit(0);
}

unsigned int pow (unsigned int x, unsigned int y) {
	int pow = 1;
	for (int i = 0; i < y; i++) {
		pow*=x;
	}
	return pow;
}
int calc_tree_size(int depth, int numOfChildren) {
	assert(numOfChildren > 0 && depth >= 0);
	if (numOfChildren == 1) {
		return depth + 1;
	}
	return (pow(numOfChildren, depth + 1) - 1) / (numOfChildren - 1);
}

int calc_lca(int index1, int index2, int depth, int numOfChildren) {
	assert (numOfChildren > 0 && index1 >= 0 && index2 >= 0);
	int *path1 = (int*) malloc(sizeof(int) * (depth + 1));
	if (path1 == NULL) {
		perror("Out of memory\n");
		return -1;
	}
	int *path2 = (int*) malloc(sizeof(int) * (depth + 1));
	if (path2 == NULL) {
		free(path1);
		perror("Out of memory\n");
		return -1;
	}
	int pathLength1, pathLength2;
	path_from_root(index1, path1, numOfChildren, &pathLength1);
	path_from_root(index2, path2, numOfChildren, &pathLength2);

	int height = 0;
	while(path1[height] == path2[height] && height < pathLength1 && height < pathLength2) {
		height++;
	}

	int index = 0;
	for (int i = 0; i < height; i++) {
		index = INDEX_FATHER_TO_CHILD(index, numOfChildren, path1[i]);
	}

	free(path1);
	free(path2);

	return index;
	
}

void path_from_root(int index, int *path, int numOfChildren, int* pathLength) {
	int length = 0;
	int temp = index;
	while(temp != 0) {
		length++;
		temp = INDEX_CHILD_TO_FATHER(temp, numOfChildren);
	}
	*pathLength = length;
	for (int i = length - 1; i > -1; i--) {
		path[i] = (index - 1) % numOfChildren;
		index = INDEX_CHILD_TO_FATHER(index, numOfChildren);
	}
}

int calc_path(int indexSrc, int indexDst, int depth, int numOfChildren, int* path, int pathLength) {
	assert(indexSrc >= 0 && indexDst >= 0 && depth > 0 && numOfChildren > 0 && path != NULL);
	int lca = calc_lca(indexSrc, indexDst, depth, numOfChildren);
	if (lca < 0) {
		perror("Out of memory\n");
		return -1;
	}
	int i = 0;
	while (indexSrc != lca && i < pathLength) {
		path[i++] = indexSrc;
		indexSrc = INDEX_CHILD_TO_FATHER(indexSrc, numOfChildren);
	} 

	int *pathDst = (int*) malloc(sizeof(int) * (depth+1));
	if (pathDst == NULL) {
		perror("Out of memory\n");
		return -1;
	}
	int pathDstLength;
	path_from_root(indexDst, pathDst, numOfChildren, &pathDstLength);
	int index = 0;
	int j = 0;
	while (index != lca) {
		index = INDEX_FATHER_TO_CHILD(index, numOfChildren, pathDst[j++]);
	}
	while(index != indexDst && i < pathLength) {
		path[i++] = index;
		index = INDEX_FATHER_TO_CHILD(index, numOfChildren, pathDst[j++]);
	} 
	if (i < pathLength) {
		path[i++] = index;
	} else {
		return -1;
		errno = ENOMEM;
	}
	free(pathDst);
	return i;
}

int main(int argc, const char* argv[]) {
	if (argc < 3) {
		perror("too few parameters\n");
	}
	int lcaMatch = 1;
	int pathMatch = 1;
	int corenerMatch = 1;
	int retVal1, retVal2;
	int depth = atoi(argv[1]);
	int numOfChildren = atoi(argv[2]);
	int treeSize = calc_tree_size(depth, numOfChildren);
	int maxPathLength =  2 * (depth + 1) - 1;
	pid_t* pathHW1 =  (pid_t*) malloc(sizeof(pid_t) * maxPathLength);
	int retVal;
	if (pathHW1 == NULL) {
		perror("Out of memory\n");
		return -1;
	}
	int* pathRef =  (int*) malloc(sizeof(int) * maxPathLength);
	if (pathRef == NULL) {
		free(pathHW1);
		perror("Out of memory\n");
		return -1;
	}
	
	 /* Declare shared memory variables */
	key_t keyPA;
	int shmidPA;

	key_t keyDone;
	int shmidDone;

	/* Initialize Shared Memory */
	keyPA = ftok(SHARED_MEM_FILE,SHARED_MEM_ID);
	printf ("key: %d\n", keyPA);
	shmidPA = shmget(keyPA, treeSize * sizeof(pid_t), 0666 | IPC_CREAT);
	printf ("shmid: %d\n", shmidPA);

	/* Initialize Shared Memory */
	keyDone = ftok(SHARED_MEM_FILE_1,SHARED_MEM_ID_1);
	printf ("key1: %d\n", keyDone);
	shmidDone = shmget(keyDone, sizeof(int), 0666 | IPC_CREAT);
	printf ("shmid1: %d\n", shmidDone);
	
	/* Attach to Shared Memory */
	processArr = shmat(shmidPA, (void *)0, 0);
	if(processArr == (pid_t *)(-1)) {
		perror("Could not attach to shared memory");
	}

	/* Attach to Shared Memory */
	isDone = shmat(shmidDone, (void *)0, 0);
	if(isDone == (int *)(-1)) {
		perror("Could not attach to shared memory");
	}

	*isDone = 0;
	
	build_process_tree(depth, numOfChildren);
	int isContinue = 0;
	while (!isContinue) {
		sleep(SLEEP_SEC);
		fflush(stdout);
		isContinue = 1;
		for (int i = 0; i < treeSize; i++) {
			if (processArr[i] == 0) {
				isContinue = 0;
			}
		}
	}
	fflush(stdout);
	int lcaHW1, lcaRef;
	pid_t pid1, pid2;
	int i;
	//Validity test, all input is well formed (pids exist, arrays allocated etc.)
	for (int index1 = 0; index1 < treeSize; index1++) {
		for (int index2 = 0; index2 < treeSize; index2++) {
			pid1 = processArr[index1];
			pid2 = processArr[index2];
			lcaHW1 = get_common_ancestor(pid1, pid2);
			lcaRef = calc_lca(index1, index2,depth, numOfChildren);
			if (lcaRef < 0) {
				free(pathHW1);
				free(pathRef);
				free(processArr);
				perror("Out of memory\n");
				return lcaRef;
			}
			//printf("LCA REF: INDEX1 = %d, , INDEX2 = %d, LCA = %d\n",index1, index2, lcaRef);
			//printf("LCA HW1: PID1 = %d, PID2 = %d, INDEX1 = %d, , INDEX2 = %d, LCA = %d\n",pid1, pid2, index1, index2, lcaHW1);
			if (lcaHW1 != processArr[lcaRef]) {
				printf("LCA Mismatch: PID1 =%d , PID2 =%d , INDEX1 = %d, , INDEX2 = %d\n",pid1, pid2,index1, index2);
				lcaMatch = 0;
			}

			retVal1 = get_path(pid1, pid2, pathHW1, maxPathLength);
			retVal2 = calc_path(index1, index2, depth, numOfChildren, pathRef, maxPathLength);
			i = 0;
			assert(retVal2 > 0);
			if (retVal1 != retVal2) {
				printf("Path Mismatch: PID1 =%d , PID2 =%d , INDEX1 = %d, , INDEX2 = %d. get_path failed.\n",pid1, pid2,index1, index2);
				pathMatch = 0;
			}
			while (pathRef[i] != index2) {
				if (processArr[pathRef[i]] != pathHW1[i++]) {
					printf("Path Mismatch: PID1 =%d , PID2 =%d , INDEX1 = %d, , INDEX2 = %d\n",pid1, pid2,index1, index2);
					pathMatch = 0;
				}
			} 
			if (processArr[pathRef[i]] != pathHW1[i]) {
				printf("Path Mismatch: PID1 =%d , PID2 =%d , INDEX1 = %d, , INDEX2 = %d\n",pid1, pid2,index1, index2);
				pathMatch = 0;
			}
			fflush(stdout);
		}
	} 

	/*Corner cases*/
	//pid1 invalid
	retVal1 = get_common_ancestor(-1, 1);
	if (retVal1 != -1 || errno != ESRCH) {
		corenerMatch = 0;
		printf ("LCA error: should return EINVAL\n");
	}
	retVal1 = get_path(-1, 1, pathHW1, maxPathLength);
	if (retVal1 != -1 || errno != ESRCH) {
		corenerMatch = 0;
		printf ("LCA error: should return EINVAL\n");
	}
	//pid2 invalid
	retVal1 = get_common_ancestor(1, -1);
	if (retVal1 != -1 || errno != ESRCH) {
		corenerMatch = 0;
		printf ("LCA error: should return EINVAL\n");
	}
	retVal1 = get_path(1, -1, pathHW1, maxPathLength);
	if (retVal1 != -1 || errno != ESRCH) {
		corenerMatch = 0;
		printf ("LCA error: should return EINVAL\n");
	}
	//both pid invalid
	retVal1 = get_common_ancestor(-1, -1);
	if (retVal1 != -1 || errno != ESRCH) {
		corenerMatch = 0;
		printf ("LCA error: should return EINVAL\n");
	}
	retVal1 = get_path(-1, -1, pathHW1, maxPathLength);
	if (retVal1 != -1 || errno != ESRCH) {
		corenerMatch = 0;
		printf ("LCA error: should return EINVAL\n");
	}
	//too short path array
	int* path = malloc(sizeof(int));
	if (path == NULL) {
		free(pathHW1);
		free(pathRef);
		perror("Out of memory\n");
		return -1;
	}
	retVal1 = get_path(1, processArr[treeSize - 1],path, 1);
	if (retVal1 != -1 || errno != ENOMEM) {
		corenerMatch = 0;
		printf ("get_path error: should return ENOMEM\n");
	}
	free(path);

	//path array doesn't belong to process address space
	retVal1 = get_path(1, 1, (int*) 0xfffffff9, 6);
	if (retVal1 != -1 || errno != EINVAL) {
		corenerMatch = 0;
		printf ("get_path error: should return EINVAL\n");
	}
	//size negative or zero
	retVal1 = get_path(1, 1, pathHW1, 0);
	if (retVal1 != -1 || errno != EINVAL) {
		corenerMatch = 0;
		printf ("get_path error: should return EINVAL\n");
	}
	retVal1 = get_path(1, 1, pathHW1, -5);
	if (retVal1 != -1 || errno != EINVAL) {
		corenerMatch = 0;
		printf ("get_path error: should return EINVAL\n");
	}
	//array null
	retVal1 = get_path(1, 1, (int*) NULL, 5);
	if (retVal1 != -1 || errno != EINVAL) {
		corenerMatch = 0;
		printf ("get_path error: should return EINVAL\n");
	}
	//check errno value order
	retVal1 = get_path(-1, 1, (int*) NULL, 5);
	if (retVal1 != -1 || errno != ESRCH) {
		corenerMatch = 0;
		printf ("get_path error: should return EINVAL\n");
	}

	retVal1 = get_path(1, processArr[treeSize - 1],NULL, 1);
	if (retVal1 != -1 || errno != EINVAL) {
		corenerMatch = 0;
		printf ("get_path error: should return ENOMEM\n");
	}


	if (lcaMatch == 1) {
		printf("get_common_ancestor: Full Tree... Success\n");
	} else {
		printf("get_common_ancestor: Full Tree... Fail\n");
	}
	if (pathMatch == 1) {
		printf("get_path: Full Tree... Success\n");
	} else {
		printf("get_path: Full Tree... Fail\n");
	}

	if (corenerMatch == 1) {
		printf("Corner cases... Success\n");
	} else {
		printf("Corner cases ... Fail\n");
	}

	*isDone = 1;
	int status;
	while(wait(&status) != -1) {}

	/* Detach from Shared Memory */
	shmdt(processArr);

	/* Detach from Shared Memory */
	shmdt(isDone);

	free(pathHW1);
	free(pathRef);
}
/************************************************************************/
